import base64
import json
import cv2
from ultralytics import YOLO
import numpy as np
from flask_cors import CORS
from flask import Flask, request, Response
import logging

# 加载YOLOv8模型
model = None
# 加载本地摄像头
camera = cv2.VideoCapture(0)


# 初始化
def create_app():
    global model
    app = Flask(__name__)
    logger = logging.getLogger(__name__)
    logger.info('加载app成功')
    CORS(app)
    # 初始化模型
    model = YOLO(model="./model/yolov8n-pose.pt")
    print('模型加载成功')
    # 加载路由和其它初始化过程
    # 载入测试图片
    test_img = cv2.imread('./model/wbb.jpg')
    print('加载初始化图片')
    # 模型预测
    results = model(test_img)
    print('进行初始预测')
    frame = results[0].plot()
    print('初始化预测成功')
    print("结果" + str(type(frame)))
    return app


# 创建应用
app = create_app()


# 本机摄像头调用
def gen_frames():
    while True:
        success, frame = camera.read()
        if not success:
            break
        else:
            # YOLOv8预测
            results = model(frame)
            # 绘制预测结果
            frame = results[0].plot()
            # Encode as a jpeg image and return
            ret, buffer = cv2.imencode('.jpg', frame)
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')


# 网络接口
@app.route('/receive_image', methods=['POST'])
def receive_image():
    # 从request中获取参数
    data = request.data.decode('utf-8')
    json_data = json.loads(data)
    str_image = json_data.get("imgData")
    img = base64.b64decode(str_image)
    img_np = np.frombuffer(img, dtype='uint8')
    new_img_np = cv2.imdecode(img_np, 1)
    # YOLOv8预测
    results = model(new_img_np)
    # # 绘制预测结果
    print("预测结果"+str(type(results[0])))
    frame = results[0].plot()
    ret, buffer = cv2.imencode('.jpg', frame)
    frame = buffer.tobytes()

    return Response(frame,
                    mimetype='multipart/x-mixed-replace; boundary=frame')


# 本机测试接口
@app.route('/video_feed')
def video_feed():
    return Response(gen_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')


if __name__ == "__main__":
    app.run(debug=True)
